import re
import os
import argparse


def parse_ascend_cpp(file_path):
    with open(file_path, 'r') as f:
        content = f.read()

    # 1. 提取类名
    class_name = None
    class_match = re.search(r'class\s+(\w+)\s*{', content)
    if class_match:
        class_name = class_match.group(1)
        print(f"class_name = {class_name}")

    # 2. 提取 __global__ __aicore__ 函数名和参数
    global_func = None
    pattern = r'extern\s+"C"\s+__global__\s+__aicore__\s+\w+\s+(\w+)\s*\(([^)]*)\)'
    global_func_match = re.search(pattern, content)

    if global_func_match:
        function_name = global_func_match.group(1)
        params_str = global_func_match.group(2)
        params = []
        if params_str.strip():
            param_parts = params_str.split(',')
            for part in param_parts:
                part = part.strip()
                words = part.split()
                if words:
                    params.append(words[-1])
        global_func = {
            "name": function_name,
            "params": params
        }
        print(f"global_func = {global_func}")

    # 3. 提取最后一个void函数
    void_pattern = r'void\s+(\w+)\s*\(([^)]*)\)\s*\{'
    void_matches = list(re.finditer(void_pattern, content))
    if void_matches:
        last_match = void_matches[-1]
        function_name = last_match.group(1)
        params_str = last_match.group(2)
        print(f"function_name = {function_name}, params_str={params_str}")
        # 解析所有参数
        all_params = []
        if params_str.strip():
            # 这里使用更复杂的解析方式来处理指针和其他类型
            param_parts = []
            # 考虑逗号内的复杂情况，比如函数指针参数
            current_part = ""
            paren_level = 0

            for char in params_str:
                if char == ',' and paren_level == 0:
                    param_parts.append(current_part.strip())
                    current_part = ""
                else:
                    current_part += char
                    if char == '(':
                        paren_level += 1
                    elif char == ')':
                        paren_level -= 1

            if current_part.strip():
                param_parts.append(current_part.strip())

            # 解析每个参数部分
            for part in param_parts:
                all_params.append(part.strip())

        void_func = {
            "name": function_name,
            "params": all_params
        }

    if not class_name:
        raise ValueError("类名未在文件中找到！")
    if not global_func_match:
        raise ValueError("__global__ __aicore__ 函数未在文件中找到！")
    if not void_matches:
        raise ValueError("void函数未在文件中找到！")

    return {
        "class_name": class_name,
        "global_func": global_func,
        "void_func": void_func
    }


def replace_and_save_cpp(src_path, dest_path, data):
    """替换类名、函数名并保存新文件"""
    with open(src_path, 'r') as f:
        content = f.read()

    # 替换类名
    content = re.sub(
        rf'\b{data["class_name"]}\b',
        'CurrOp',
        content
    )

    # 替换全局函数名
    content = re.sub(
        rf'\b{data["global_func"]["name"]}\b',
        'curr_op',
        content
    )

    # 替换全局函数名
    content = re.sub(
        rf'\b{data["void_func"]["name"]}\b',
        'curr_op_do',
        content
    )

    # 保存新文件
    with open(dest_path, 'w') as f:
        f.write(content)
    print(f"生成替换后的 CPP 文件: {dest_path}")


def generate_pybind_code(last_void_func_params, output_path):
    """生成 Pybind11 适配代码"""
    # 提取参数类型和名称（如 "uint8_t* x" 提取为 "uint8_t" 和 "*x"）
    param_types = [p.split()[0] for p in last_void_func_params]
    param_names = [p.split()[-1] for p in last_void_func_params]
    # 处理星号*的位置，比如把 type 和 name 分别为"uint8_t" 和 "*x"的情况修改成"uint8_t*" 和 "x"
    for i, param_name in enumerate(param_names):
        if param_name[0] == '*':
            param_names[i] = param_name[1:]
            param_types[i] = param_types[i] + "*"
        else:
            param_names[i] = param_name

    # 生成 extern 声明
    extern_decl = f"extern void curr_op_do({', '.join(last_void_func_params)});"

    # 生成函数参数部分
    func_params = ',\n  '.join([f'const at::Tensor &{name}' for name in param_names[2:]])  # 跳过 blockDim 和 stream
    call_params_list = []
    for i in range(2, len(param_names)):
        print(f"param_types[{i}] = {param_types[i]}")
        if param_types[i] == "uint8_t*":
            param_types[i] = "const_cast<void *>"
        else:
            param_types[i] = f"reinterpret_cast<{param_types[i]}>"
        call_params_list.append(f'{param_types[i]} ({param_names[i]}.data_ptr())')
    call_params = ', '.join(call_params_list)
    print(f"call_params = {call_params}")

    pybind_code = f"""#include <iostream>
#include <pybind11/pybind11.h>
#include "curr_op_tiling.h"
#include "aclrtlaunch_curr_op.h"
#include <torch/extension.h>
#include "torch_npu/csrc/core/npu/NPUStream.h"
#include "data_utils.h"
#include "custom_tiling.h"
#include "acl/acl.h"

{extern_decl}

namespace my_op {{

void run_curr_op(
  {func_params})
{{
  auto aclStream = c10_npu::getCurrentNPUStream().stream(true);
  ACLRT_LAUNCH_KERNEL(curr_op)(blockDim, aclStream, {call_params});
}}

}}  // namespace my_op

PYBIND11_MODULE(curr_op_py, m) {{
  m.doc() = "curr_op pybind11 interfaces";
  m.def("run_curr_op", &my_op::run_curr_op, "");
}}
"""
    with open(output_path, 'w') as f:
        f.write(pybind_code)
    print(f"生成 Pybind11 文件: {output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='处理 Ascend C 文件生成适配代码')
    parser.add_argument('input', help='输入的 Ascend C 文件路径')
    parser.add_argument('--cpp-output', default='modified.cpp', help='修改后的 CPP 文件路径')
    parser.add_argument('--pybind-output', default='curr_op_pybind.cpp', help='Pybind11 文件路径')
    args = parser.parse_args()

    # 解析原始文件
    try:
        data = parse_ascend_cpp(args.input)
        # 1. 生成替换后的 CPP 文件
        replace_and_save_cpp(args.input, args.cpp_output, data)
        # 2. 生成 Pybind11 代码（使用最后一个 void 函数的参数）
        generate_pybind_code(data["void_func"]["params"], args.pybind_output)
    except ValueError as e:
        print(f"错误: {e}")
